package com.baishancloud.log.common.sink

import cn.hutool.http.HttpUtil
import cn.hutool.json.JSONUtil
import org.apache.flink.api.common.state.{ListState, ListStateDescriptor}
import org.apache.flink.api.common.typeinfo.{TypeHint, TypeInformation}
import org.apache.flink.configuration.Configuration
import org.apache.flink.runtime.state.{FunctionInitializationContext, FunctionSnapshotContext}
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction
import org.apache.flink.streaming.api.functions.sink.{RichSinkFunction, SinkFunction}
import org.slf4j.{Logger, LoggerFactory}

import java.util.{Timer, TimerTask}
import scala.collection.JavaConverters.mutableSeqAsJavaListConverter
import scala.collection.mutable
import scala.util.control.Breaks

/**
 * @param outThreshold 写入mallard阈值，默认为int最大值
 * @param outInterval  写入mallard时间间隔，默认5秒，单位：秒
 * @author ziqiang.wang
 * @date 2021/11/16 09:53
 */
class MallardSink(outThreshold: Int = Int.MaxValue, outInterval: Long = 5) extends RichSinkFunction[String] with CheckpointedFunction {

  private val LOGGER: Logger = LoggerFactory.getLogger(classOf[MallardSink])

  //当前系统时间
  var currentTimeMillis: Long = System.currentTimeMillis()
  var lastOutTimeMillis: Long = currentTimeMillis - outInterval * 1000
  //算子内数据结构，存放累计的结果数据
  val results: mutable.ListBuffer[String] = mutable.ListBuffer[String]()
  //将未被发送的数据放到该state，以便checkpoint时进行持久化
  var resultsState: ListState[String] = _
  val url: String = "http://127.0.0.1:10699/v2/push"

  override def open(parameters: Configuration): Unit = {
    //定时更新
    new Timer("currentTimeMillis", true).schedule(new TimerTask {
      override def run(): Unit = currentTimeMillis = System.currentTimeMillis()
    }, 1000, 1000)
  }


  override def invoke(value: String, context: SinkFunction.Context): Unit = {
    results.+=(value)
    if ((currentTimeMillis - lastOutTimeMillis >= outInterval * 1000 || results.size >= outThreshold) && results.nonEmpty) {
      // 新数据到达时的当前时间达到输出时间间隔，或者是累计数据数量到达设定阈值，则进行post请求发送数据
      val body: StringBuilder = new StringBuilder()
      body.append("[")
      results.foreach((result: String) => body.append(result).append(","))
      body.deleteCharAt(body.length - 1).append("]")
      val httpBody: String = body.toString()
      Breaks.breakable {
        for (i <- 1 to 3) {
          val response: String = HttpUtil.post(url, httpBody)
          if (response.contains("OK")) {
            results.clear()
            Breaks.break()
            lastOutTimeMillis = currentTimeMillis
          } else {
            LOGGER.error(s"post请求错误，进行第${i}次重试，错误信息：\n$response,body数据：\n${JSONUtil.toJsonStr(results.asJava)}")
          }
        }
      }
    }
  }


  /**
   * 做快照时执行，清空状态内数据，然后将算子内数据添加到状态
   */
  override def snapshotState(context: FunctionSnapshotContext): Unit = {
    resultsState.clear()
    results.foreach(r => resultsState.add(r))
  }

  /**
   * 初始化checkpoint，将状态内数据恢复到算子内
   */
  override def initializeState(context: FunctionInitializationContext): Unit = {
    val resultStateDescriptor: ListStateDescriptor[String] = new ListStateDescriptor[String]("上批数据缓存", TypeInformation.of(new TypeHint[String] {}))
    resultsState = context.getOperatorStateStore.getListState(resultStateDescriptor)
    if (context.isRestored) {
      if (resultsState.get() != null) {
        resultsState.get().forEach(result => {
          results.+=(result)
        })
      }
    }
  }


}
